import numpy as np
import config

class SubChannel:
    def __init__(self, n_axial=None, n_channels=None):
        """初始化子通道
        
        Args:
            n_axial: 轴向网格数，如果为None则从配置文件读取
            n_channels: 子通道数，如果为None则从配置文件读取
        """
        # 从配置文件或参数获取网格数
        self.n_axial = n_axial if n_axial is not None else config.GRID_PARAMS['n_axial']
        self.n_channels = n_channels if n_channels is not None else config.GRID_PARAMS['n_channels']
        
        # 初始化横向连接关系
        self.lateral_connections = config.LATERAL_CONNECTION_PARAMS['connections']
        self.heated_boundaries = config.LATERAL_CONNECTION_PARAMS['heated_boundaries']
        
        # 几何参数
        self.dx = config.GEOMETRY_PARAMS['channel_length'] / self.n_axial
        self.dy = config.GEOMETRY_PARAMS['channel_width']
        self.hydraulic_diameter = np.array([config.GEOMETRY_PARAMS['hydraulic_diameter']] * self.n_channels)  # 水力直径
        self.flow_area = np.array([config.GEOMETRY_PARAMS['flow_area']] * self.n_channels)  # 流通面积
        self.wetted_perimeter = np.array([config.GEOMETRY_PARAMS['wetted_perimeter']] * self.n_channels)  # 湿周
        self.heated_perimeter = np.array([config.GEOMETRY_PARAMS['heated_perimeter']] * self.n_channels)  # 加热周长
        
        # 流动场
        self.pressure = np.zeros((self.n_axial, self.n_channels))  # 压力场
        self.void_fraction = np.zeros((self.n_axial, self.n_channels))  # 空泡分数
        self.vel_liquid = np.zeros((self.n_axial, self.n_channels))  # 液相速度
        self.vel_vapor = np.zeros((self.n_axial, self.n_channels))  # 气相速度
        self.lateral_velocity = np.zeros((self.n_axial, self.n_channels))  # 横向速度
        
        # 温度场
        self.temp_liquid = np.zeros((self.n_axial, self.n_channels))  # 液相温度
        self.temp_vapor = np.zeros((self.n_axial, self.n_channels))  # 气相温度
        self.temp_wall = np.zeros((self.n_axial, self.n_channels))  # 壁面温度
        
        # 物性场 - 基本物性
        self.rho_liquid = np.zeros((self.n_axial, self.n_channels))  # 液相密度
        self.rho_vapor = np.zeros((self.n_axial, self.n_channels))  # 气相密度
        self.mu_liquid = np.zeros((self.n_axial, self.n_channels))  # 液相粘度
        self.mu_vapor = np.zeros((self.n_axial, self.n_channels))  # 气相粘度
        self.k_liquid = np.zeros((self.n_axial, self.n_channels))  # 液相导热系数
        self.k_vapor = np.zeros((self.n_axial, self.n_channels))  # 气相导热系数
        self.cp_liquid = np.zeros((self.n_axial, self.n_channels))  # 液相比热容
        self.cp_vapor = np.zeros((self.n_axial, self.n_channels))  # 气相比热容
        
        # 物性场 - 焓值
        self.enthalpy_liquid = np.zeros((self.n_axial, self.n_channels))  # 液相焓值
        self.enthalpy_vapor = np.zeros((self.n_axial, self.n_channels))  # 气相焓值
        self.enthalpy_mixture = np.zeros((self.n_axial, self.n_channels))  # 混合焓值
        
        # 物性场 - 相变和界面物性
        self.latent_heat = np.zeros((self.n_axial, self.n_channels))  # 汽化潜热
        self.surface_tension = np.zeros((self.n_axial, self.n_channels))  # 表面张力
        
        # 初始化物性场
        self._initialize_properties()
        
    def _initialize_properties(self):
        """初始化物性场"""
        # 设置初始压力分布（考虑重力引起的静压差）
        P_inlet = config.INITIAL_CONDITIONS['inlet_pressure']
        T_inlet = config.INITIAL_CONDITIONS['inlet_temperature']
        rho_l = self._calculate_liquid_density(P_inlet, T_inlet)
        
        # 设置初始压力分布
        for i in range(self.n_axial):
            dP_gravity = rho_l * config.PHYSICAL_PARAMS['gravity'] * i * self.dx
            self.pressure[i,:] = P_inlet - dP_gravity
            
        # 设置初始温度
        self.temp_liquid[:,:] = T_inlet
        self.temp_vapor[:,:] = T_inlet
        self.temp_wall[:,:] = T_inlet + 5  # 假设初始壁温比流体温度高5K
        
        # 设置初始空泡分数
        self.void_fraction[:,:] = config.INITIAL_CONDITIONS.get('inlet_void_fraction', 0.0)
        
        # 设置初始速度
        G_inlet = config.INITIAL_CONDITIONS.get('inlet_mass_flux', 1000.0)  # kg/m²s
        A_flow = config.GEOMETRY_PARAMS['flow_area']
        
        # 计算初始速度
        for j in range(self.n_channels):
            alpha = self.void_fraction[0,j]
            rho_l = self._calculate_liquid_density(P_inlet, T_inlet)
            rho_v = self._calculate_vapor_density(P_inlet, T_inlet)
            
            # 计算相速度
            if alpha < 1e-6:  # 纯液相
                self.vel_liquid[:, j] = G_inlet / rho_l
                self.vel_vapor[:, j] = 0.0
            elif alpha > 0.999:  # 纯气相
                self.vel_liquid[:, j] = 0.0
                self.vel_vapor[:, j] = G_inlet / rho_v
            else:  # 两相流
                x = alpha * rho_v / (alpha * rho_v + (1 - alpha) * rho_l)  # 质量含气率
                self.vel_liquid[:, j] = G_inlet * (1 - x) / (rho_l * (1 - alpha))
                self.vel_vapor[:, j] = G_inlet * x / (rho_v * alpha)
        
        # 更新所有物性
        self.update_fluid_properties()
        
    def update_fluid_properties(self):
        """更新流体物性"""
        for i in range(self.n_axial):
            for j in range(self.n_channels):
                P = self.pressure[i,j]
                T_l = self.temp_liquid[i,j]
                T_v = self.temp_vapor[i,j]
                
                # 更新基本物性
                self.rho_liquid[i,j] = self._calculate_liquid_density(P, T_l)
                self.rho_vapor[i,j] = self._calculate_vapor_density(P, T_v)
                self.mu_liquid[i,j] = self._calculate_liquid_viscosity(T_l)
                self.mu_vapor[i,j] = self._calculate_vapor_viscosity(T_v)
                self.k_liquid[i,j] = self._calculate_liquid_conductivity(T_l)
                self.k_vapor[i,j] = self._calculate_vapor_conductivity(T_v)
                self.cp_liquid[i,j] = self._calculate_liquid_specific_heat(T_l)
                self.cp_vapor[i,j] = self._calculate_vapor_specific_heat(T_v)
                
                # 更新焓值
                self.enthalpy_liquid[i,j] = self._calculate_liquid_enthalpy(T_l)
                self.enthalpy_vapor[i,j] = self._calculate_vapor_enthalpy(P, T_v)
                self.enthalpy_mixture[i,j] = self._calculate_mixture_enthalpy(i, j)
                
    def _calculate_liquid_density(self, P, T):
        """计算液相密度 (kg/m³)"""
        # 使用简化的水密度计算公式
        T_C = T - 273.15  # 转换为摄氏度
        return 1000 * (1 - ((T_C - 3.98)**2) / 50000)
        
    def _calculate_vapor_density(self, P, T):
        """计算气相密度 (kg/m³)"""
        # 使用理想气体状态方程
        R = 461.5  # 水蒸气气体常数 J/(kg·K)
        return P / (R * T)
        
    def _calculate_liquid_viscosity(self, T):
        """计算液相粘度 (Pa·s)"""
        # 使用简化的水粘度计算公式
        T_C = T - 273.15
        return 0.001 * (1 + 0.0337 * T_C + 0.000221 * T_C**2) ** (-1)
        
    def _calculate_vapor_viscosity(self, T):
        """计算气相粘度 (Pa·s)"""
        # 使用Sutherland公式
        mu_0 = 1.85e-5  # 参考粘度
        T_0 = 373.15  # 参考温度
        S = 620  # Sutherland常数
        return mu_0 * (T/T_0)**1.5 * (T_0 + S)/(T + S)
        
    def _calculate_liquid_conductivity(self, T):
        """计算液相导热系数 (W/(m·K))"""
        # 使用简化的水导热系数计算公式
        T_C = T - 273.15
        return 0.6 - 0.0005 * T_C
        
    def _calculate_vapor_conductivity(self, T):
        """计算气相导热系数 (W/(m·K))"""
        # 使用简化的水蒸气导热系数计算公式
        return 0.016 + 5e-5 * (T - 273.15)
        
    def _calculate_liquid_specific_heat(self, T):
        """计算液相比热容 (J/(kg·K))"""
        # 使用简化的水比热容计算公式
        T_C = T - 273.15
        return 4200 * (1 + 0.0002 * T_C)
        
    def _calculate_vapor_specific_heat(self, T):
        """计算气相比热容 (J/(kg·K))"""
        # 使用简化的水蒸气比热容计算公式
        return 2000 + T - 273.15
        
    def _calculate_liquid_enthalpy(self, T):
        """计算液焓值 (J/kg)"""
        # 以0℃液态水为基准
        return self._calculate_liquid_specific_heat(T) * (T - 273.15)
        
    def _calculate_vapor_enthalpy(self, P, T):
        """计算气相焓值 (J/kg)"""
        # 计算饱和温度
        T_sat = self._calculate_saturation_temperature(P)
        
        # 以0℃液态水为基准，加上汽化潜热和过热焓
        h_fg = self._calculate_latent_heat(T_sat)  # 汽化潜热
        h_f = self._calculate_liquid_enthalpy(T_sat)  # 饱和水焓值
        cp_v = self._calculate_vapor_specific_heat(T)  # 气相比热
        
        return h_f + h_fg + cp_v * (T - T_sat)
        
    def _calculate_mixture_enthalpy(self, i, j):
        """计算混合焓值 (J/kg)"""
        alpha = self.void_fraction[i,j]
        rho_l = self.rho_liquid[i,j]
        rho_v = self.rho_vapor[i,j]
        h_l = self.enthalpy_liquid[i,j]
        h_v = self.enthalpy_vapor[i,j]
        
        # 计算质量含气率
        x = alpha * rho_v / (alpha * rho_v + (1 - alpha) * rho_l)
        
        # 计算混合焓值
        return (1 - x) * h_l + x * h_v
        
    def _calculate_saturation_temperature(self, P):
        """计算饱和温度 (K)"""
        # 使用简化的饱和温度计算公式
        P_MPa = P / 1e6  # 转换为MPa
        return 273.15 + 100 * P_MPa / 0.101325
        
    def _calculate_latent_heat(self, T):
        """计算汽化潜热 (J/kg)"""
        # 使用Watson公式
        T_c = 647.15  # 临界温度 (K)
        T_r = min(max(T / T_c, 0.0), 1.0)  # 约化温度，限制在[0,1]范围内
        h_fg_0 = 2257e3  # 在1个大气压下的汽化潜热 (J/kg)
        
        if T_r >= 1.0:  # 超临界状态
            return 0.0
        else:
            return h_fg_0 * (1 - T_r)**0.38
        
    def get_lateral_connection_type(self, channel_idx, neighbor_idx):
        """获取两个通道之间的连接类型
        
        Args:
            channel_idx: 当前通道索引
            neighbor_idx: 相邻通道索引
            
        Returns:
            connection_type: 连接类型 (1-子通道间连接, 2-固体边界, 3-加热边界)
            connection_width: 连接宽度
        """
        # 查找连接关系
        for conn in self.lateral_connections:
            if (conn[0] == channel_idx and conn[1] == neighbor_idx) or \
               (conn[1] == channel_idx and conn[0] == neighbor_idx):
                return conn[2], conn[3]
        return None, None
        
    def get_heated_boundary_width(self, channel_idx):
        """获取通道的加热边界宽度
        
        Args:
            channel_idx: 通道索引
            
        Returns:
            heated_width: 加热边界宽度
        """
        for boundary in self.heated_boundaries:
            if boundary[0] == channel_idx:
                return boundary[1]
        return 0.0
        
    def get_connected_channels(self, channel_idx):
        """获取与指定通道相连的所有通道
        
        Args:
            channel_idx: 通道索引
            
        Returns:
            connected_channels: 相连通道列表，每个元素为(通道索引, 连接类型, 连接宽度)
        """
        connected = []
        for conn in self.lateral_connections:
            if conn[0] == channel_idx:
                connected.append((conn[1], conn[2], conn[3]))
            elif conn[1] == channel_idx:
                connected.append((conn[0], conn[2], conn[3]))
        return connected